src/twinkle/infra/_ray/resource_manager.py (198-199)

The logic for determining gpu_ranks_local has been refactored. Instead of directly indexing visible_device_list, it now iterates through worker_ranks and uses self.visible_devices[node_rank][r % nproc_per_node]. This change is necessary because self.visible_devices is now a list of lists (list of devices per node), whereas visible_device_list was likely a flattened list of all visible devices. This is a critical fix for correctly mapping GPU ranks to their local device IDs on specific nodes.

                        gpu_ranks_local = []
                        for r in worker_ranks:
                            node_rank = r // nproc_per_node
                            node_ranks.append(node_rank)
                            gpu_ranks = self.visible_devices[node_rank][r % nproc_per_node]
                            gpu_ranks_local.append(gpu_ranks)

src/twinkle/infra/_ray/resource_manager.py (214)

Similar to the previous comment, the change from visible_device_list[alloc_rank % nproc_per_node] to self.visible_devices[node_rank][alloc_rank % nproc_per_node] is crucial for correctly indexing the visible devices based on the node rank. This ensures that the correct GPU is assigned to the worker, reflecting the new structure of self.visible_devices.

                        gpu_rank = self.visible_devices[node_rank][alloc_rank % nproc_per_node]

src/twinkle/utils/init.py (10)

The import of GPU, NPU, Platform, ensure_hccl_socket_env, and ensure_npu_backend is now from .platforms. This is a significant refactoring, consolidating all platform-specific logic into a new platforms submodule. This improves the organization of the utils package by grouping related functionalities.

from .platforms import GPU, NPU, Platform, ensure_hccl_socket_env, ensure_npu_backend

src/twinkle/utils/init.py (12)

The import of pad_sequence_to_length, selective_log_softmax, and stateless_init_process_group is now from .torch_utils. This indicates that stateless_init_process_group, previously in network.py, has been moved to torch_utils.py, suggesting it's more closely related to PyTorch utilities. This improves the logical grouping of functions.

from .torch_utils import pad_sequence_to_length, selective_log_softmax, stateless_init_process_group, to_device

src/twinkle/checkpoint_engine/base.py (102)

Similar to the previous change, using 'torch.Tensor' as a string literal for type hinting is appropriate here, especially since torch is now conditionally imported for type checking. This maintains type correctness without requiring a runtime import of torch if it's not otherwise used in the module.

    async def send_weights(self, weights: Generator[tuple[str, 'torch.Tensor'], None, None]):

src/twinkle/checkpoint_engine/base.py (115)

Consistent application of string literal type hints for torch.Tensor is good for the same reasons mentioned above.

    async def receive_weights(self) -> AsyncGenerator[tuple[str, 'torch.Tensor'], None]:

src/twinkle/infra/_ray/resource_manager.py (176)

The removal of assert len(groups) == len(visible_devices) indicates a change in how visible_devices is expected to be used or initialized. This assertion might have been too restrictive or incorrect given the new logic. Ensure that the new logic correctly handles cases where len(groups) and len(visible_devices) might differ, or that visible_devices is now structured differently.

        for group in groups:

src/twinkle/checkpoint_engine/base.py (3)

The torch import is moved inside a TYPE_CHECKING block. This is a good practice for type hinting imports that are not strictly needed at runtime, which can help reduce import times and avoid circular dependencies in some cases. However, ensure that torch is always imported at runtime if it's used outside of type hints in the module.

from abc import ABC, abstractmethod

src/twinkle/advantage/init.py (34-35)

These entries in __all__ are no longer valid since the corresponding functions have been removed. Removing them ensures that the module's public API accurately reflects its contents.

]

src/twinkle/model/megatron/model/gpt_bridge.py (21-22)

The import of is_last_rank has been moved from twinkle.utils.platform to twinkle.utils. This aligns with the overall refactoring of platform-related utilities into a more centralized location within the twinkle.utils package. This change improves consistency and simplifies the import structure.

from twinkle.utils import (MxFp4Dequantizer, SafetensorLazyLoader, StreamingSafetensorSaver, deep_getattr, get_logger,
                           get_modules_to_not_convert, get_multimodal_target_regex, is_last_rank, requires)

src/twinkle/advantage/init.py (7-27)

The removal of these temporary helper functions (compute_advantages and compute_advantages_rloo) is a good step towards cleaning up the codebase. It aligns with the TODO comment indicating they were temporary and not suitable for production. This improves the clarity and long-term maintainability of the advantage module by removing deprecated code.

src/twinkle/checkpoint_engine/base.py (11-12)

Using string literals for type hints like 'torch.Size' and 'torch.dtype' when torch is only imported for type checking is the correct approach. This prevents a runtime import error if torch is not available or if it causes a circular dependency.

src/twinkle/utils/network.py (1-43)

The removal of os, torch, datetime, and the HCCL-related environment variable definitions and functions (_HCCL_IF_BASE_PORT_ENV, _HCCL_HOST_SOCKET_PORT_RANGE_ENV, _HCCL_NPU_SOCKET_PORT_RANGE_ENV, _derive_hccl_socket_env_defaults, _ensure_hccl_socket_env) from network.py is a key part of the refactoring. These functionalities have been moved to the new platforms submodule, specifically npu.py, and stateless_init_process_group to torch_utils.py. This makes network.py more focused on generic network utilities.

import socket
from typing import Optional

src/twinkle/utils/network.py (90-170)

The stateless_init_process_group function has been moved from network.py to torch_utils.py. This is a logical relocation as this function is specifically designed for PyTorch distributed process group initialization, even if it uses network concepts. This improves the thematic organization of the utils package.

src/twinkle/utils/platforms/init.py (1-4)

This new __init__.py file in twinkle/utils/platforms serves to expose the Platform, GPU, MPS, NPU classes and related functions (is_mps_available, ensure_hccl_socket_env, ensure_npu_backend). This is a crucial part of the refactoring, creating a dedicated package for platform-specific abstractions and utilities. This improves modularity and makes the codebase easier to navigate and extend for different hardware platforms.

src/twinkle/utils/platforms/base.py (1-138)

This new file base.py within twinkle/utils/platforms defines the abstract Platform class and its static methods. This centralizes common platform-agnostic functionalities and provides a clear interface for platform-specific implementations (GPU, NPU, MPS). This is a significant improvement in code organization and extensibility.

src/twinkle/utils/platforms/gpu.py (1-26)

This new file gpu.py within twinkle/utils/platforms defines the GPU class, inheriting from Platform. It encapsulates GPU-specific environment variables, device prefixes, and device retrieval logic. This is a good example of adhering to the new modular structure for platform-specific implementations.

src/twinkle/utils/platforms/mps.py (1-42)

This new file mps.py within twinkle/utils/platforms defines the MPS class and the is_mps_available function. It encapsulates MPS-specific environment variables, device prefixes, and device retrieval logic. This is a good example of adhering to the new modular structure for platform-specific implementations.

src/twinkle/utils/platforms/npu.py (1-135)

This new file npu.py within twinkle/utils/platforms defines the NPU class and related helper functions for Ascend NPU. It centralizes HCCL environment variable handling, NPU device ID resolution, and the ensure_npu_backend function. This is a crucial part of the refactoring, moving all NPU-specific logic into a dedicated module and improving code organization.

src/twinkle/utils/torch_utils.py (112-194)

The stateless_init_process_group function has been moved from network.py to torch_utils.py. This is a logical relocation as this function is specifically designed for PyTorch distributed process group initialization, even if it uses network concepts. This improves the thematic organization of the utils package. Additionally, the import for ensure_hccl_socket_env is now correctly pointing to twinkle.utils.platforms.